package cn.monkey.family.server.domain.uc.auth.token;


import cn.monkey.family.data.context.RequestContext;
import cn.monkey.family.data.http.HttpHeaderConstants;
import jakarta.servlet.*;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.http.HttpHeaders;

import java.io.IOException;
import java.util.Optional;
import java.util.function.Consumer;

public class DownloadTemporaryTokenFilter implements Filter {

    private String temporaryTokenName = "temporary-token";

    private final RequestContextRepository requestContextRepository;

    public DownloadTemporaryTokenFilter(RequestContextRepository requestContextRepository) {
        this.requestContextRepository = requestContextRepository;
    }

    public void setTemporaryTokenName(String temporaryTokenName) {
        this.temporaryTokenName = temporaryTokenName;
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
        if (request instanceof HttpServletRequest httpServletRequest
                && response instanceof HttpServletResponse httpServletResponse) {
            this.doFilter0(httpServletRequest, httpServletResponse, chain);
            return;
        }
        chain.doFilter(request, response);
    }

    static class HeaderAddingRequestWrapper extends HttpServletRequestWrapper {

        private final HttpHeaders headers;

        HeaderAddingRequestWrapper(HttpServletRequest httpServletRequest,
                                   HttpHeaders headers) {
            super(httpServletRequest);
            this.headers = headers;
        }


        @Override
        public String getHeader(String name) {
            if (this.headers.containsKey(name)) {
                return this.headers.getFirst(name);
            }
            return super.getHeader(name);
        }
    }

    static <T> void tryAdd(Consumer<T> consumer, T value) {
        if (value != null) {
            consumer.accept(value);
        }
    }


    private void doFilter0(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, FilterChain chain) throws ServletException, IOException {
        String temporaryToken = httpServletRequest.getParameter(this.temporaryTokenName);
        if (temporaryToken == null) {
            chain.doFilter(httpServletRequest, httpServletResponse);
            return;
        }
        RequestContext requestContext = this.requestContextRepository.get(temporaryToken);
        if (requestContext == null) {
            chain.doFilter(httpServletRequest, httpServletResponse);
            return;
        }
        HttpHeaders headers = new HttpHeaders();
        Optional.ofNullable(requestContext.getToken()).ifPresent(s -> headers.add(HttpHeaderConstants.AUTHORIZATION_KEY, s));
        Optional.ofNullable(requestContext.getOrgId()).ifPresent(s -> headers.add(HttpHeaderConstants.ORG_ID_KEY, s));
        Optional.ofNullable(requestContext.getPlatformId()).ifPresent(s -> headers.add(HttpHeaderConstants.PLATFORM_ID_KEY, s));
        Optional.ofNullable(requestContext.getTraceId()).ifPresent(s -> headers.add(HttpHeaderConstants.TRACE_ID_KEY, s));
        chain.doFilter(new HeaderAddingRequestWrapper(httpServletRequest, headers), httpServletResponse);
    }
}
